import cv2 
import numpy as np
import os
import shutil


# 切割图片
def split_img(img_path):
    img = cv2.imread(img_path)
    img = np.asarray(img)
    
    width = img.shape[1]
    cv2.imwrite(img_path[0:-3] + '_0.png', img[:,:width // 2])
    cv2.imwrite(img_path[0:-3] + '_1.png', img[:,width // 2 :])
# split_img(r'C:\Users\vanlance\Desktop\20160629002549.bmp')


# ROI
def get_img_roi(img_path):
    img = cv2.imread(img_path)
    img = np.asarray(img)
    
    new_img = img[300:900, 500:1500]
    cv2.imwrite('crop.jpg', new_img)
# get_img_roi(r'C:\Users\vanlance\Desktop\2\2016\out\392.jpg')


# 删除长宽比特定的图片
def delete_img_for_the_ratio(img_path):
    for p in os.listdir(img_path):
        file_ = img_path + os.sep + p
        img = cv2.imread(file_)
        img = np.asarray(img)
        h,w = img.shape[:2]
        if 1.4 < float(h) / w < 18 and h < 1200:
            os.remove(file_)
        elif 1.4 < float(w) / h < 18 and w < 1200:
            os.remove(file_)
# delete_img_for_the_ratio(r'C:\Users\vanlance\Desktop\2\2016\out')

def resize_img(img_path):
    new_path = '_.'.join(img_path.split('.'))
    img = cv2.imread(img_path,0)
    height,width = img.shape
    save_img = cv2.resize(img,(width // 2,height // 2))
    cv2.imwrite(new_path, save_img)


# 视频转图片
def video2img(dir_path):
    if os.path.isfile(dir_path):
        img_path = os.path.join(os.path.dirname(dir_path), 'imgs')
        os.makedirs(img_path,exist_ok=True) 

        
        video_path = dir_path

        cap = cv2.VideoCapture(video_path)
        i = 0
        while 1:
            ret, frame = cap.read()
            if ret:
                # cv2.imshow("capture", frame)
                if (i+5) % 5 == 0:
                    cv2.imwrite(os.path.join(img_path, str(i)+ '.png'), frame)  #[246:936,512:1275]   [137:1019,540:1453]
                i += 1
                
                # if cv2.waitKey(1) & 0xFF == ord('c'):
                #     cv2.imwrite(str(i) + '.png',frame)
                   
                if cv2.waitKey(1) & 0xFF == ord('b'):
                    break
            else: break
    else :
        img_path = os.path.join(dir_path, 'imgs')
        os.makedirs(img_path,exist_ok=True)

        for name in os.listdir(dir_path):
            
            video_path = os.path.join(dir_path, name)
            if os.path.isdir(video_path):
                continue
            cap = cv2.VideoCapture(video_path)

            print(video_path)
            i = 0
            while 1:
                ret, frame = cap.read()
                if ret:
                    # cv2.imshow("capture", frame)
                    if i%15==0:
                        p = os.path.join(img_path, name[-8:-4] + '_' + str(i)+ '.png')
                        # print(p)
                        cv2.imwrite(p, frame)
                    i += 1
                    
                    # if cv2.waitKey(1) & 0xFF == ord('c'):
                    #     cv2.imwrite(str(i) + '.png',frame)
                    
                    if cv2.waitKey(1) & 0xFF == ord('b'):
                        break
                else: break

    
                
            cap.release()
            cv2.destroyAllWindows()


def resize_img(path):
    img = cv2.imread(path,0)
    img1 = cv2.resize(img, (300,60))
    new_path = path.replace("dataset",'dataset1')
    cv2.imwrite(new_path,img1)

def main(path):
    
    for p in os.listdir(path):
        file_path = os.path.join(path,p)
        resize_img(file_path)

def rename(path):
    files = [os.path.join(path,p) for p in os.listdir(path)]
    imgs=[]
    for num,f in enumerate(files):
        img = cv2.imread(f)
        imgs.append(img)
        os.remove(f)

    for n,i in enumerate(imgs):
        cv2.imwrite(os.path.join(path, str(n)+ ".png"), i)

def shuffle(path):
    files = [os.path.join(path, p) for p in os.listdir(path)]
    imgs=[]
    for num,f in enumerate(files):
        img = cv2.imread(f,0)
        imgs.append(img)

    label_name = os.path.basename(path)
    train_path = path.replace('dataset1','dataset')
    train_path = os.path.join(os.path.dirname(train_path),'train')
    test_path = train_path.replace('train', 'test')
    
    np.random.shuffle(imgs)
    length = len(imgs)
    for num,img in enumerate(imgs):
        if num < length * 0.25:
            save_path = os.path.join(test_path,label_name)
        else:
            save_path = os.path.join(train_path,label_name)


        os.makedirs(save_path,exist_ok=True)
        cv2.imwrite(os.path.join(save_path, str(num) + '.png' ) ,img)

def make_label(path):
    with open(os.path.join(path , 'label.txt'),'a') as f:
        labels = os.listdir(path)
        for l in labels:
            new_path = os.path.join(path, l)
            if os.path.isdir(new_path):
                for p in os.listdir(new_path):
                    f.write(l + os.sep + p + " " + l + '\n')

def random_choice(path, num):
    file_names = [os.path.join(path, p) for p in os.listdir(path)]
    np.random.shuffle(file_names)

    length = len(file_names)
    if length < num:
        num = length

    for name in file_names[:num]:
        pass

# 上下左右翻转图片，扩充
def dilate_img(path):
    files = [os.path.join(path,p) for p in os.listdir(path)]
    all_img = []
    for f in files:
        if os.path.isfile(f):
            img = cv2.imread(f,0)
            img1 = img[:,::-1]
            img2 = img[::-1,:]
            img3 = img2[:,::-1]
            all_img.extend([img,img1,img3, img2])
            os.remove(f)
    np.random.shuffle(all_img)

    for num,image in enumerate(all_img):
        write_path = path +os.sep + str(num) + ".png"
        cv2.imwrite(write_path, image)

def move_file(sub_file_paths, test_path):
    np.random.shuffle(sub_file_paths)
    np.random.shuffle(sub_file_paths)
    length = len(sub_file_paths)
    std_num = length // 6

    for _, image in enumerate(sub_file_paths[:std_num]):
        shutil.move(image, test_path)


# 递归将训练数据分到测试文件夹
def make_train_val_set(train_path, exclude=['label.txt']):
    test_path = train_path.replace('/train/', '/test/')
    if not os.path.exists(test_path):
        os.mkdir(test_path)

    sub_paths = [os.path.join(train_path, p) for p in os.listdir(train_path) ]
    sub_file_paths = []
    sub_dir_paths = []
    [sub_file_paths.append(sfp) if os.path.isfile(sfp) else sub_dir_paths.append(sfp) for sfp in sub_paths ]
    if len(sub_file_paths) > 0:
        move_file([p for p in sub_file_paths for e in exclude if e not in p ], test_path)
    if len(sub_dir_paths) > 0:
        for sub_dir_path in sub_dir_paths:
            make_train_val_set(sub_dir_path)


if __name__ == '__main__':
    # path = r'C:\Users\vanlance\Documents\Visual Studio 2013\Projects\RecogDLL_split_img\x64\Release'
    # for p in os.listdir(path):
    # split_img(os.path.join(path, p))
    # resize_img(r'C:\Users\vanlance\Desktop\git\swls_file_transfer\002.jpg')
    # video2img(r'C:\Users\vanlance\Desktop\2\video')
    dilate_img('/home/swls/work_dir/git/socket_recog/train_model/dataset/train/1')
    dilate_img('/home/swls/work_dir/git/socket_recog/train_model/dataset/train/0')
    make_train_val_set('/home/swls/work_dir/git/socket_recog/train_model/dataset/train/')
    make_label('/home/swls/work_dir/git/socket_recog/train_model/dataset/train/')
    make_label('/home/swls/work_dir/git/socket_recog/train_model/dataset/test/')